package com.wfh.aigenerator.ai;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.wfh.aigenerator.ai.guardrail.PromptSafetyInputGuardrail;
import com.wfh.aigenerator.ai.guardrail.RetryOutputGuardrail;
import com.wfh.aigenerator.ai.tool.*;
import com.wfh.aigenerator.common.ErrorCode;
import com.wfh.aigenerator.exception.BusinessException;
import com.wfh.aigenerator.model.entity.App;
import com.wfh.aigenerator.model.enums.CodeGenTypeEnums;
import com.wfh.aigenerator.service.AppService;
import com.wfh.aigenerator.service.ChatHistoryService;
import com.wfh.aigenerator.utils.SpringContextUtil;
import dev.langchain4j.community.store.memory.chat.redis.RedisChatMemoryStore;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.guardrail.InputGuardrail;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.service.AiServices;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Lazy;

import java.time.Duration;

import static com.wfh.aigenerator.model.enums.CodeGenTypeEnums.*;

/**
 * @Title: AiGenServiceFactory
 * @Author wangfenghuan
 * @Package com.wfh.aigenerator.ai
 * @Date 2025/7/29 20:14
 * @description:
 */
@Configuration
@Slf4j
public class AiGenServiceFactory {

    @Resource(name = "openAiChatModel")
    private ChatModel chatModel;


    @Resource
    @Lazy
    private AppService appService;

    @Resource
    private ToolManager toolManager;


    @Resource
    private RedisChatMemoryStore redisChatMemoryStore;


    @Resource
    private ChatHistoryService chatHistoryService;

    /**
     * AI 服务实例缓存
     * 缓存策略：
     * - 最大缓存 1000 个实例
     * - 写入后 30 分钟过期
     * - 访问后 10 分钟过期
     */
    private final Cache<String, AiGenService> serviceCache = Caffeine.newBuilder()
            .maximumSize(1000)
            .expireAfterWrite(Duration.ofMinutes(30))
            .expireAfterAccess(Duration.ofMinutes(10))
            .removalListener((key, value, cause) -> {
                log.debug("AI 服务实例被移除，appId: {}, 原因: {}", key, cause);
            })
            .build();


    /**
     * 根据appId获取ai服务示例，这里是为了兼容历史的逻辑（携带缓存）
     * @param appId
     * @return
     */
    public AiGenService getAiGenService(Long appId) {
        App app = appService.getById(appId);
        if (app == null){
            throw new BusinessException(ErrorCode.OPERATION_ERROR);
        }
        String codeGenType = app.getCodeGenType();
        CodeGenTypeEnums byValue = getByValue(codeGenType);
        return getAiGenService(appId, byValue);
    }

    /**
     * 根据appId获取ai服务示例，携带缓存
     * @param appId
     * @param codeGenTypeEnums
     * @return
     */
    public AiGenService getAiGenService(Long appId, CodeGenTypeEnums codeGenTypeEnums) {
         String cacheKey = buildCacheKey(appId, codeGenTypeEnums);
        return serviceCache.get(cacheKey, key -> createAiGenService(appId, codeGenTypeEnums));
    }


    private AiGenService createAiGenService(long appId, CodeGenTypeEnums codeGenType) {
        // 注意：这里不要提前构建 chatMemory，而是延迟构建
        return switch (codeGenType) {
            case VUE -> {
                StreamingChatModel thinkingStreamChatModelPrototype = SpringContextUtil.getBean("thinkingStreamChatModelPrototype", StreamingChatModel.class);
                yield  AiServices.builder(AiGenService.class)
                        .streamingChatModel(thinkingStreamChatModelPrototype)
                        // ✅ 动态创建
                        .chatMemoryProvider(memoryId -> createChatMemory(appId))
                        .tools(toolManager.getAllTools())
                        .tools(new ImageTool())
                        .hallucinatedToolNameStrategy(toolExecutionRequest ->
                                ToolExecutionResultMessage.from(toolExecutionRequest,
                                        "Error: there is no tool called " + toolExecutionRequest.name()))
                        .inputGuardrails(new PromptSafetyInputGuardrail())
                        // .outputGuardrails(new RetryOutputGuardrail())
                        .maxSequentialToolsInvocations(20)
                        .build();
            }

            case HTML, MULTI_FILE -> {
                StreamingChatModel streamingChatModelPrototype = SpringContextUtil.getBean("streamingChatModelPrototype", StreamingChatModel.class);
                yield  AiServices.builder(AiGenService.class)
                    .chatModel(chatModel)
                        .tools(new ImageTool())
                    .streamingChatModel(streamingChatModelPrototype)
                    // ✅ 统一使用 provider
                    .chatMemoryProvider(memoryId -> createChatMemory(appId))
                        .inputGuardrails(new PromptSafetyInputGuardrail())
                        .outputGuardrails(new RetryOutputGuardrail())
                    .build();
            }

            default -> throw new BusinessException(ErrorCode.SYSTEM_ERROR,
                    "不支持的代码生成类型: " + codeGenType.getValue());
        };
    }

    /**
     * 创建 ChatMemory 实例
     * @param appId
     * @return
     */
    private ChatMemory createChatMemory(Long appId) {
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .id(appId)
                .chatMemoryStore(redisChatMemoryStore)
                .maxMessages(20)
                .build();

        // 加载历史记录
        chatHistoryService.loadToMemory(appId, chatMemory, 20);

        return chatMemory;
    }


    /**
     * 构造缓存key
     * @param appId
     * @param codeGenTypeEnums
     * @return
     */
    private String buildCacheKey(Long appId, CodeGenTypeEnums codeGenTypeEnums) {
        return appId + "-" + codeGenTypeEnums.name();
    }
}
